#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
编排器组件，负责协调多智能体之间的消息传递和流程管理
支持通过 MemorySystem 管理状态，并处理 Web 应用的交互请求
"""

from email import message
import logging
import uuid
import json
import re
import os
from datetime import datetime, timezone  # Make sure timezone is imported
from typing import Dict, List, Any, Optional, Tuple

try:
    from agents.base_agent import BaseAgent
    from agents.companion import CompanionAgent
    from utils.memory_system import MemorySystem
except ImportError as e:
    logging.critical(f"Orchestrator 无法导入核心依赖: {e}", exc_info=True)

    class PlaceholderAgent:
        pass

    class PlaceholderMemorySystem:
        pass

    BaseAgent = PlaceholderAgent  # type: ignore
    CompanionAgent = PlaceholderAgent  # type: ignore
    MemorySystem = PlaceholderMemorySystem  # type: ignore
    STAGE_INTENT = STAGE_ASKING_IDENTITY = STAGE_GUIDING = STAGE_CHATTING = STAGE_CHAT_ENDED = STAGE_CLARIFICATION = STAGE_ERROR = STAGE_COMPLETED = "unknown_stage"

logger = logging.getLogger("AI-MindCare-System-Orchestrator")

DEFAULT_DEPARTMENT = "心理树洞"


class Orchestrator:
    """编排器，负责协调多智能体系统的交互流程"""

    def __init__(self,
                 companion: Optional[CompanionAgent] = None,
                 memory_system: Optional[MemorySystem] = None):

        if not isinstance(companion, CompanionAgent):
            raise TypeError("必须提供有效的 CompanionAgent 实例。")
        if not isinstance(memory_system, MemorySystem):
            raise TypeError("必须提供有效的 MemorySystem 实例。")

        self.companion = companion
        self.memory_system = memory_system
        self.agents: Dict[str, BaseAgent] = {"companion": companion}
        logger.info("编排器初始化完成")

    def _log_interaction_event(self, user_id: str, conversation_id: str,
                               event_details: Dict[str, Any]):
        """记录交互事件到会话上下文，确保日志包含足够信息在前端显示，并保存到文件"""
        if not user_id or not conversation_id or not self.memory_system:
            return

        try:
            # 获取当前上下文
            context = self.memory_system.get_conversation_context(
                user_id, conversation_id)
            if not context:
                logger.warning(f"无法记录交互事件: 找不到上下文 {conversation_id}")
                return

            # 添加时间戳timestamp
            event_details["timestamp"] = datetime.now(
                timezone.utc).strftime("%H:%M:%S.%f")[:-3] + "Z"

            # 确保详情字段存在，并且至少包含基本信息
            event_details.setdefault("details", {})

            # 为不同类型的交互添加更丰富的信息
            source = event_details.get("source", "")
            target = event_details.get("target", "")
            action = event_details.get("action", "")

            # 根据不同的交互类型，丰富详情信息
            if "发送请求" in action or "触发" in action:
                # 添加请求类型说明
                if "content_keys" in event_details.get("details", {}):
                    keys = event_details["details"]["content_keys"]
                    if isinstance(keys, list) and len(keys) > 0:
                        event_details["details"]["request_type"] = keys[0]

            elif "收到响应" in action or "收到反馈结果" in action:
                # 添加状态描述
                status = event_details.get("details", {}).get("status")
                if status:
                    event_details["details"][
                        "status_description"] = self._get_status_description(
                            status)

            elif "用户输入" in action:
                # 保持消息片段
                pass

            elif "错误" in action:
                # 确保错误详情可见
                pass

            # 添加适合人类阅读的描述
            event_details[
                "human_readable"] = self._generate_human_readable_description(
                    source, target, action, event_details.get("details", {}))
            # 记录日志项
            log_msg = f"Event Logged [{conversation_id}]: {source} -> {target} ({action})"
            if 'status' in event_details.get('details', {}):
                log_msg += f" Status: {event_details['details']['status']}"
            logger.debug(log_msg)

            # 保存交互日志到文件
            self._save_interaction_log_to_file(conversation_id, event_details)

        except Exception as e:
            logger.error(
                f"Failed to log interaction event for {conversation_id}: {e}",
                exc_info=True)

    def _save_interaction_log_to_file(self, conversation_id: str,
                                      event_details: Dict[str, Any]):
        """保存交互日志到文件，以便在会话结束后仍然可以查看"""
        try:
            # 日志文件路径，以conversation_id命名
            log_dir = "logs"
            if not os.path.exists(log_dir):
                os.makedirs(log_dir)

            log_file_path = os.path.join(
                log_dir, f"interaction_log_{conversation_id}.json")

            # 准备要写入的数据
            log_entry = {
                "timestamp":
                event_details.get(
                    "timestamp",
                    datetime.now(timezone.utc).strftime("%H:%M:%S.%f")[:-3] +
                    "Z"),
                "source":
                event_details.get("source", "unknown"),
                "target":
                event_details.get("target", "unknown"),
                "action":
                event_details.get("action", "unknown"),
                "details":
                event_details.get("details", {}),
                "human_readable":
                event_details.get("human_readable", "")
            }

            # 读取现有日志文件（如果存在）
            existing_logs = []
            if os.path.exists(log_file_path):
                try:
                    with open(log_file_path, 'r', encoding='utf-8') as f:
                        existing_logs = json.load(f)
                except json.JSONDecodeError:
                    logger.warning(f"无法解析现有日志文件 {log_file_path}，将创建新文件")

            # 追加新日志
            existing_logs.append(log_entry)

            # 写回文件
            with open(log_file_path, 'w', encoding='utf-8') as f:
                json.dump(existing_logs, f, ensure_ascii=False, indent=2)

            logger.debug(f"已将交互事件保存到文件 {log_file_path}")

        except Exception as e:
            logger.error(f"保存交互日志到文件时出错: {e}", exc_info=True)

    def _get_status_description(self, status: str) -> str:
        """根据状态码返回人类可读的状态描述"""
        status_descriptions = {
            # 心理陪伴智能体状态
            "in_conversation": "对话中",
            "conversation_complete": "对话结束",

            # 接待状态
            "in_progress": "接待进行中",
            "completed": "接待完成",

            # 通用状态
            "error": "出现错误",
            "success": "成功",
            "unhandled": "未处理",
        }
        return status_descriptions.get(status, status)

    def _generate_human_readable_description(self, source: str, target: str,
                                             action: str,
                                             details: Dict[str, Any]) -> str:
        """生成人类可读的事件描述"""
        # 初始信息
        description = ""

        # 处理不同类型的源
        if source == "User":
            description = "用户发送消息"
        elif source == "System":
            description = "系统"
        elif source == "Orchestrator":
            if "触发" in action:
                if target == "Companion":
                    description = "系统将患者转给咨询师进行诊断"
                else:
                    description = f"系统触发{target}服务"
            else:
                description = f"系统向{target}发送请求"
        elif source == "receptionist":
            description = "接待员完成接待流程"
        elif source == "companion":
            if action == "收到诊断触发响应":
                description = "心理咨询师开始诊断"
            elif "诊断" in action:
                description = "心理咨询师完成诊断"
            else:
                description = "心理咨询师回复"

        # 添加详情
        if "error" in details:
            description += f"（错误：{details['error']}）"
        elif "message" in details:
            description += f"（{details['message']}）"
        elif "status_description" in details:
            description += f"（{details['status_description']}）"

        return description

    def start_new_web_conversation(self) -> Tuple[str, str]:
        conversation_id = f"web_{uuid.uuid4()}"  # 系统生成对话ID
        logger.info(f"Orchestrator: 开始新的 Web 对话，ID: {conversation_id}")
        initial_message = "有烦恼吗？找我聊聊，帮你赶走坏心情😊"  # 这个后期可以更改，暂时保留这个选项
        return conversation_id, initial_message

    def handle_web_message(
            self,
            user_id: str,
            conversation_id: str,
            user_message: str,  # 用户最新的消息
            user_buffer: Dict[str, Any],
            long_term_memory) -> Dict[str, Any]:
        logger.info(
            f"Orchestrator: 处理 Web 消息 (Conversation: {conversation_id}): '{user_message[:50]}...'"
        )

        if not all(agent in self.agents
                   for agent in ["companion"]) or not self.memory_system:
            logger.error(
                f"Orchestrator 缺少核心组件 (Conversation: {conversation_id})")
            return {
                "status": "error",
                "message": "系统内部错误，核心服务不可用。",
                "agent_type": "system"
            }

        target_agent: Optional[BaseAgent] = None
        agent_input_content: Optional[Dict[str, Any]] = None
        request_key_for_log: str = "unknown_request"

        if user_buffer:  # 如果会话缓冲区不为空
            # --- Log User Input ---
            self._log_interaction_event(
                user_id, conversation_id, {
                    "source": "User",
                    "target": "System",
                    "action": "用户输入",
                    "details": {
                        "message_snippet": user_message[:30] + "..."
                    }
                })
            target_agent = self.agents.get("companion")
            agent_input_content = {
                "user_response": {
                    "user_id": user_id,
                    "answer": user_message,
                    "context_id": conversation_id,
                    "is_followup": True,  # 新增字段标识后续对话，followup跟踪
                }
            }
            request_key_for_log = "user_response"

        else:  # No active context - First message
            logger.info(f"初始化新咨询会话 {conversation_id}，直接转心理陪伴智能体")

            # 构造咨询师初始请求
            target_agent = self.companion
            agent_input_content = {
                "conversation_request": {  # 使用companion智能体的标准请求格式
                    "user_id": user_id,
                    "conversation_id": conversation_id,
                    "raw_message": user_message  # 保留原始消息
                }
            }

            # Log user input now that we have user_id
            self._log_interaction_event(
                user_id,
                conversation_id,
                {
                    "source": "User",
                    "target": "System",
                    "action": "用户输入 (首次)",  # 标识这是首次输入
                    "details": {
                        "message_snippet": user_message[:30] + "..."
                    }
                })

            request_key_for_log = "conversation_request"

        if not target_agent or not agent_input_content:
            logger.error(
                f"无法确定处理消息的目标 Agent 或输入内容 (Conversation: {conversation_id})")
            if conversation_id:
                self.memory_system.delete_conversation_context(conversation_id)
            return {
                "status": "error",
                "message": "系统内部错误，无法处理您的请求。",
                "agent_type": "system"
            }

        try:
            logger.info(
                f"将消息转发给 Agent: {target_agent.role} (Name: {target_agent.name}), Conversation: {conversation_id}"
            )
            self._log_interaction_event(
                user_id=user_id,
                conversation_id=conversation_id,
                event_details={
                    "source": "Orchestrator",
                    "target": target_agent.role,
                    "action": f"发送请求 ({request_key_for_log})",
                    "details": {
                        "content_keys": list(agent_input_content.keys())
                    }
                })
            agent_message_envelope = {
                "sender_id": "orchestrator_web_handler",
                "sender_role": "orchestrator",
                "content": agent_input_content,
                "conversation_id": conversation_id,
                "user_id": user_id,
                "long_term_memory": long_term_memory,
                "messages": user_buffer  # 传递完整的消息缓冲区，实质上是简化后的上下文
            }
            agent_response_msg = target_agent.receive_message(
                agent_message_envelope)
            resp_content_keys = list(
                agent_response_msg.get("content", {}).keys())
            resp_status = agent_response_msg.get("content",
                                                 {}).get("status", "unknown")
            self._log_interaction_event(
                user_id, conversation_id, {
                    "source": target_agent.role,
                    "target": "Orchestrator",
                    "action": "收到响应",
                    "details": {
                        "status": resp_status,
                        "content_keys": resp_content_keys
                    }
                })
            return self._process_agent_response_for_web(
                user_id, conversation_id, agent_response_msg)
        except Exception as e:
            logger.error(
                f"与 Agent ({target_agent.role}) 交互时出错 (Conversation: {conversation_id}): {e}",
                exc_info=True)
            self._log_interaction_event(
                user_id, conversation_id, {
                    "source": "Orchestrator",
                    "target": target_agent.role,
                    "action": "交互错误",
                    "details": {
                        "error": str(e)
                    }
                })
            if conversation_id:
                self.memory_system.delete_conversation_context(conversation_id)
            return {
                "status": "error",
                "message": f"系统在与 {target_agent.role} 沟通时出现问题，请稍后重试。",
                "agent_type": "system"
            }

    # --- END REINSTATED HELPER METHODS ---

    def _process_agent_response_for_web(
        self,
        user_id: str,
        conversation_id: str,
        agent_response: Dict[str, Any],
    ) -> Dict[str, Any]:
        if not agent_response or not isinstance(agent_response, dict):
            logger.error(
                f"收到无效的 Agent 响应 (Conversation: {conversation_id}): {agent_response}"
            )
            return {
                "status": "error",
                "message": "系统内部通信错误。",
                "agent_type": "system"
            }

        agent_role = agent_response.get("sender_role", "system")
        content = agent_response.get("content", {})
        status_from_agent = content.get("status", "unknown")
        next_stage_from_agent = content.get("next_stage")

        # Ensure the message field is always a valid string
        default_message = f"来自 {agent_role} 的响应处理中..."
        message_from_agent = content.get("message", default_message)
        if message_from_agent is None or not isinstance(
                message_from_agent, str):
            message_from_agent = default_message
            logger.warning(
                f"Agent返回了非字符串消息 (Conversation: {conversation_id}, Role: {agent_role}), 使用默认消息"
            )

        web_response = {
            "agent_type": agent_role,
            "message": message_from_agent,
            "status": "in_progress",
            "conversation_id": conversation_id,
            # "prescription": None,
            # "diagnosis_info": None,
            # "context_id": conversation_id
        }

        logger.info(
            f"处理来自 {agent_role} 的响应 (AgentStatus: {status_from_agent}, NextStage: {next_stage_from_agent}), Content Keys: {list(content.keys())}"
        )

        if agent_role == "companion":
            context_id_from_companion = content.get("context_id",
                                                    conversation_id)
            web_response["context_id"] = context_id_from_companion
            if status_from_agent == "in_conversation":
                companion_question = content.get("question")
                if not companion_question or not isinstance(
                        companion_question, str):
                    companion_question = "心理咨询师需了解更多信息"
                web_response["message"] = companion_question
                web_response["status"] = "in_progress"

            else:
                error_msg = content.get("message", "心理咨询师处理时出错")
                if not isinstance(error_msg, str):
                    error_msg = "心理咨询师处理时出错"
                web_response["message"] = error_msg
                web_response["status"] = "error"
                self._log_interaction_event(
                    user_id, conversation_id, {
                        "source": "Orchestrator",
                        "target": "System",
                        "action": "错误",
                        "details": {
                            "message": f"心理咨询师返回错误/未知状态: {status_from_agent}"
                        }
                    })
        elif status_from_agent == 'unhandled':
            error_msg = content.get("message", "系统无法处理请求")
            if not isinstance(error_msg, str):
                error_msg = "系统无法处理请求"
            web_response["message"] = error_msg
            web_response["status"] = "error"
            self._log_interaction_event(
                user_id, conversation_id, {
                    "source": "Orchestrator",
                    "target": "System",
                    "action": "错误",
                    "details": {
                        "message": f"{agent_role}无法处理请求"
                    }
                })

        elif status_from_agent == "error":
            error_msg = content.get("message", "处理时未知错误")
            if not isinstance(error_msg, str):
                error_msg = "处理时未知错误"
            web_response["message"] = error_msg
            web_response["status"] = "error"
            self._log_interaction_event(
                user_id, conversation_id, {
                    "source": "Orchestrator",
                    "target": "System",
                    "action": "错误",
                    "details": {
                        "message": f"{agent_role}报告错误: {error_msg}"
                    }
                })

        final_context = self.memory_system.get_conversation_context(
            user_id, conversation_id)
        web_response["interaction_log"] = final_context.get(
            "interaction_log", []) if final_context else [{
                "timestamp":
                datetime.now(timezone.utc).strftime("%H:%M:%S.%f")[:-3] + "Z",
                "source":
                "System",
                "target":
                "",
                "action":
                "错误",
                "details":
                "无法检索日志"
            }]

        terminal_status_web = [
            "completed_no_prescription", "completed_prescription",
            "error_prescription_rejected", "error"
        ]
        if web_response["status"] in terminal_status_web:
            logger.info(
                f"Web 咨询 {conversation_id} 结束: {web_response['status']}")
            self._log_interaction_event(
                user_id, conversation_id, {
                    "source": "System",
                    "target": "",
                    "action": "咨询结束",
                    "details": {
                        "status": web_response['status']
                    }
                })

        logger.info(
            f"Orchestrator: 返回给 Web (Conversation: {conversation_id}): Status={web_response['status']}, Agent={web_response['agent_type']}, Msg='{web_response['message'][:70]}...'"
        )
        return web_response
